package com.hliushi.spark.sql

import org.apache.spark.sql.types._
import org.apache.spark.sql.{RelationalGroupedDataset, SparkSession}
import org.junit.Test

/**
 * descriptions:
 *
 * author: Hliushi
 * date: 2021/5/19 20:48
 */
class AggProcessor {

  // 1.创建SparkSession
  private val spark: SparkSession = SparkSession.builder()
    .appName("agg_processor")
    .master("local[6]")
    .getOrCreate()

  import org.apache.spark.sql.functions._
  import spark.implicits._


  /**
   * 需求: 统计每个月PM值的平均数
   */
  @Test
  def groupBy(): Unit = {
    // 2.数据读取
    //   2.1指定Schema, 不要自动推断
    val schema = StructType(
      List(
        StructField("id", LongType),
        StructField("year", IntegerType),
        StructField("month", IntegerType),
        StructField("day", IntegerType),
        StructField("hour", IntegerType),
        StructField("season", IntegerType),
        StructField("pm", DoubleType)
      )
    )

    val sourceDF = spark.read
      .schema(schema)
      .option("header", value = true)
      .csv("dataset/beijingpm_with_nan.csv")

    // 3.数据去掉空值
    val cleanDF = sourceDF.where($"pm" =!= Double.NaN)

    // 分组
    // def groupBy(cols: Column*): RelationalGroupedDataset
    val groupedDF: RelationalGroupedDataset = cleanDF.groupBy($"year", $"month")

    // 4.使用functions函数来完成聚合
    import org.apache.spark.sql.functions._

    // 本质上, avg这个函数定义了一个操作, 把表达式设置给pm列
    // select avg(pm) from ... group by
    groupedDF.agg(avg($"pm").as("pm_avg"))
      .orderBy($"pm_avg".desc)
      .show()

    // 5.使用GroupedDataset的API来完成聚合
    /**
     * select($"avg(pm)".as("pm_avg")) 重新修改列名, 和下面的效果是一样的
     * withColumnRenamed("avg(pm)", "pm_avg")
     */
    groupedDF.avg("pm")
      .withColumnRenamed("avg(pm)", "pm_avg")
      //.select($"avg(pm)".as("pm_avg"))
      .orderBy($"pm_avg".desc)
      .show()

    spark.stop()
  }


  /**
   * 需求1: 不同来源PM的统计
   * 需求2: 在同一个月,不同来源的PM值的平均值是多少
   * 需求3: 同一年,不同来源的PM值平均是多少
   * 需求4: 整体上来看,不同来源的PM值是多少
   */
  @Test
  def multiAgg(): Unit = {
    // 2.数据读取
    //   2.1指定Schema, 不要自动推断
    val schema = StructType(
      List(
        StructField("source", StringType),
        StructField("year", IntegerType),
        StructField("month", IntegerType),
        StructField("day", IntegerType),
        StructField("hour", IntegerType),
        StructField("season", IntegerType),
        StructField("pm", DoubleType)
      )
    )

    val pmFinal = spark.read
      .schema(schema)
      .option("header", value = true)
      .csv("dataset/pm_final.csv")

    //pmFinal.show()


    // 需求1:不同年, 不同来源, PM值的平均数
    // select source, year avg(pm) as pm_avg from ... group by source, year
    val postAndYearDF = pmFinal
      .groupBy($"source", $"year") // 按照source和year聚合
      .avg("pm")
      .withColumnRenamed("avg(pm)", "pm_avg")

    // 需求2:在整个数据集中, 按照不同的来源来统计PM值的平均数
    // select source, avg(pm) as pm_avg from ... group by source
    val postDF = pmFinal
      .groupBy($"source") // 按照source聚合
      .avg("pm")
      .withColumnRenamed("avg(pm)", "pm_avg")
      //  def lit(literal: Any): Column
      .select($"source", lit(null).as("year"), $"pm_avg")

    // 需求3:全局所有的计量者, 和日期的PM值的平均值   groupBy null
    val totalDF = pmFinal
      .groupBy() // 全局聚合
      .avg("pm")
      .withColumnRenamed("avg(pm)", "pm_avg")
      //  def lit(literal: Any): Column
      .select(lit(null).as("source"), lit(null).as("year"), $"pm_avg")


    // 合并在同一个数据集
    // org.apache.spark.sql.AnalysisException:
    // Union can only be performed on tables with the same number of columns, but the first table has 3 columns and the second table has 2 columns;;

    postAndYearDF.union(postDF).union(totalDF)
      .sort($"source".asc_nulls_last, $"year".asc_nulls_last, $"pm_avg").show()

    //
    //  +-------+----+------------------+
    //  | source|year|            pm_avg|
    //  +-------+----+------------------+
    //  | dongsi|2013|  93.2090724784592|
    //  | dongsi|2014| 87.08640822045773|
    //  | dongsi|2015|  87.4922056770591|
    //  | dongsi|null| 89.15443876736389|
    //  |us_post|2010|104.04572982326042|
    //  |us_post|2011|  99.0932403834184|
    //  |us_post|2012| 90.53876763535511|
    //  |us_post|2013|101.71110855035722|
    //  |us_post|2014| 97.73409537004964|
    //  |us_post|2015| 82.78472946356158|
    //  |us_post|null| 95.90424117331851|
    //  |   null|null| 93.66274738530468|
    //  +-------+----+------------------+
  }

  /**
   * 需求1: 每个城市, 每年的销售额
   * 需求2: 每个城市, 一共的销售额
   * 需求3: 总体销售额
   * #
   * rollup操作符其实就是groupBy的一个扩展, rollup会对传入的列进行滚动groupBy,
   * groupBy的次数为列数量 + 1, 最后一次是对整个数据集进行整合
   *
   * def rollup(cols: Column*): RelationalGroupedDataset
   * def rollup(col1: String, cols: String*): RelationalGroupedDataset
   */
  @Test
  def rollup(): Unit = {
    import org.apache.spark.sql.functions._

    val sales = Seq(
      ("Beijing", 2016, 100),
      ("Beijing", 2017, 200),
      ("Shanghai", 2015, 50),
      ("Shanghai", 2016, 150),
      ("Guangzhou", 2017, 50)
    ).toDF("city", "year", "amount")


    /**
     * 滚动分组, rollup(A, B)   根据A, B两列分组, 总共有AB, A, null
     */
    sales.rollup($"city", $"year")
      .agg(sum($"amount") as "sum_amount")
      .sort($"city".asc_nulls_last, $"year".asc_nulls_last)
      .show()

    //  +---------+----+----------+
    //  |     city|year|sum_amount|
    //  +---------+----+----------+
    //  |  Beijing|2016|       100|    <-- 北京2016的小计
    //  |  Beijing|2017|       200|
    //  |  Beijing|null|       300|    <-- 北京的总计
    //  |Guangzhou|2017|        50|
    //  |Guangzhou|null|        50|
    //  | Shanghai|2015|        50|
    //  | Shanghai|2016|       150|
    //  | Shanghai|null|       200|
    //  |     null|null|       550|     <-- 整个数据集的总计
    //  +---------+----+----------+
  }


  /**
   * 与上面那个 multiAgg() 方法相对应
   * multiAgg() 是使用三个 groupBy() 方法实现 rollup("source", "year")的效果
   */
  @Test
  def rollup1(): Unit = {
    // 1.读取数据集
    //   1.1指定Schema, 不要自动推断
    val schema = StructType(
      List(
        StructField("source", StringType),
        StructField("year", IntegerType),
        StructField("month", IntegerType),
        StructField("day", IntegerType),
        StructField("hour", IntegerType),
        StructField("season", IntegerType),
        StructField("pm", DoubleType)
      )
    )

    val pmFinal = spark.read
      .schema(schema)
      .option("header", value = true)
      .csv("dataset/pm_final.csv")


    // 2.聚合和统计
    //  需求1: 每个PM值计量者, 每年PM值统计的平均数    groupBy source year
    //  需求2: 每个PM值计量者, 整体上的PM平均值        groupBy source
    //  需求3: 全局所有的计量者, 和日期的PM值的平均值   groupBy null

    pmFinal.rollup($"source", $"year")
      .avg("pm")
      .withColumnRenamed("avg(pm)", "pms")
      .sort($"source".asc_nulls_last, $"year".asc_nulls_last)
      .show()
  }


  /**
   * cube的功能和rollup是一样的, 但也有区别, 区别去下
   * rollup(A, B).sum()...
   * .    其结果集中会有三种数据形式: A B, A null, null null
   * .    可以发现结果集没有对B列的聚合结果
   * cube(A, B).sum()..
   * .    其结果集中会有时钟数据形式: A B, A null, null B, null null
   * .    cube比rollup多了一个 null B, 也就是说, rollup只会按照第一列进行组合聚合,
   * .    但是cube会将全部列组合聚合
   *
   */
  @Test
  def cube(): Unit = {
    // 1.读取数据集
    //   1.1指定Schema, 不要自动推断
    val schema = StructType(
      List(
        StructField("source", StringType),
        StructField("year", IntegerType),
        StructField("month", IntegerType),
        StructField("day", IntegerType),
        StructField("hour", IntegerType),
        StructField("season", IntegerType),
        StructField("pm", DoubleType)
      )
    )

    val pmFinal = spark.read
      .schema(schema)
      .option("header", value = true)
      .csv("dataset/pm_final.csv")

    pmFinal.cube($"source", $"year")
      .avg("pm")
      .withColumnRenamed("avg(pm)", "pms")
      .sort($"source".asc_nulls_last, $"year".asc_nulls_last)
      .show()

    //  +-------+----+------------------+
    //  | source|year|               pms|
    //  +-------+----+------------------+
    //  | dongsi|2013|  93.2090724784592|
    //  | dongsi|2014| 87.08640822045773|
    //  | dongsi|2015|  87.4922056770591|
    //  | dongsi|null| 89.15443876736389|
    //  |us_post|2010|104.04572982326042|
    //  |us_post|2011|  99.0932403834184|
    //  |us_post|2012| 90.53876763535511|
    //  |us_post|2013|101.71110855035722|
    //  |us_post|2014| 97.73409537004964|
    //  |us_post|2015| 82.78472946356158|
    //  |us_post|null| 95.90424117331851|
    //  |   null|2010|104.04572982326042|
    //  |   null|2011|  99.0932403834184|
    //  |   null|2012| 90.53876763535511|
    //  |   null|2013| 97.66173808086904|
    //  |   null|2014| 92.44023222060957|
    //  |   null|2015| 85.13368549370175|
    //  |   null|null| 93.66274738530468|
    //  +-------+----+------------------+
  }


  /**
   * SparkSQL中支持SQL语句实现cube功能
   * SparkSQL支持GROUPINGS SETS语句, 可以随机排序组合控制分组聚合的顺序和组成
   * 即可以实现cube也可以实现rollup的功能
   */
  @Test
  def cubeSQL(): Unit = {
    // 1.读取数据集
    //   1.1指定Schema, 不要自动推断
    val schema = StructType(
      List(
        StructField("source", StringType),
        StructField("year", IntegerType),
        StructField("month", IntegerType),
        StructField("day", IntegerType),
        StructField("hour", IntegerType),
        StructField("season", IntegerType),
        StructField("pm", DoubleType)
      )
    )

    val pmFinal = spark.read
      .schema(schema)
      .option("header", value = true)
      .csv("dataset/pm_final.csv")

    pmFinal.createOrReplaceTempView("pm_final")

    // 使用写SQL方式
    val sql =
      """
        |select
        |   source, year, avg(pm) as pms
        |from
        |   pm_final
        |group by source, year
        |grouping sets((source, year), (source), (year), ())
        |order by source asc nulls last, year asc nulls last
        |""".stripMargin

    val result = spark.sql(sql)
    result.show()
  }
}